import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_ollama import ChatOllama
from fastapi import APIRouter, Query
from fastapi.responses import JSONResponse

router = APIRouter()
# 定义 Prompt 模板
prompt = PromptTemplate.from_template(
    "你是一个专业的律师，请你回答我提出的法律问题，并给出法律条文依据，我的问题是：{question}"
)
# 初始化大模型（Ollama）
llm = ChatOllama(
    model="qwen3:0.6b",  # 换成你在 Ollama 里有的模型名称
    base_url="http://172.16.21.38:11436"  # 服务器地址
)
parser = StrOutputParser()

@router.get("/d21prompttemplate")
async def d21prompttemplate(question: str = Query(..., description="请输入你的法律问题")):
    """接收用户输入的问题，返回大模型的回答"""
    # 生成 prompt
    prompt_value = prompt.invoke({"question": question})
    # 调用模型
    result = llm.invoke(prompt_value)
    # 解析输出
    answer = parser.invoke(result)
    return JSONResponse(content={
        "question": question,
        "prompt": prompt_value.to_string(),
        "answer": answer
    })
